'''
红色的羊和蓝色的羊竞争性的吃草。
'''
import numpy as np
import takagiabm as tak

def createSheep(color):# 调用这个函数生成一只“羊”

    if(color=='#0000ff'):
        a = tak.GridAgent()
        triangle2 = tak.scale(tak.takTriangle, 3)  # 设置形状为takagi库的三角形，缩放比例为3.
        a.setTrait('energy',10)
        a.setShape(triangle2)  # 设置形状为三角形。
        a.setSpeed(np.array([1, 0]))  # 设置初始速度
        a.setColor('#0000ff')  # 设置颜色
    elif(color=='#ff0000'):
        a = tak.GridAgent()
        triangle2 = tak.scale(tak.takTriangle, 3)  # 设置形状为takagi库的三角形，缩放比例为3.
        a.setTrait('energy',10)
        a.setShape(triangle2)  # 设置形状为三角形。
        a.setSpeed(np.array([1, 0]))  # 设置初始速度
        a.setColor('#ff0000')  # 设置颜色
    return a

def agentSetup(model:tak.GridModel):

    for i in range(5):
        a=createSheep('#0000ff')
        model.addAgent(a,randomPlace=True)
    for i in range(5):
        a=createSheep('#ff0000')
        model.addAgent(a,randomPlace=True)

    for i in range(model.width):
        for j in range(model.height):
            if(np.random.random()<0.5):
                model.grid.cells[j][i].setTrait('color','#00ff00')


def agentLoop(agent):  # 代理人(也就是羊)执行的函数
    global stat,sheepGrowth
    color = agent.getCellColor()
    if (color == '#00ff00'):
        agent.setCellColor('#ffffff')
        agent.properties['energy']+=5

    if(np.random.random()<sheepGrowth.value):
        if (agent.getTrait('energy') > 10):
            a=createSheep(agent.getTrait('color'))
            a.setTrait('energy',5)

            agent.setTrait('energy',agent.getTrait('energy')-5)

            a.setPos(agent.pos.astype(np.float))
            agent.model.addAgent(a)


    energy=agent.getTrait('energy')
    agent.setTrait('energy',energy-1)
    if(agent.getTrait('energy')==0):
        agent.model.removeAgent(agent)
        return
    agent.turn(np.random.random() * np.pi -np.pi/2)
    agent.move()

def cellStep(cell):
    global grassGrowth
    if(np.random.random()<grassGrowth.value):
        if(cell.getColor()=='#ffffff'):
            cell.setColor('#00ff00')
def cellOnClicked(cell):
    print('clicked!!!'*100)
    if(cell.getColor()=='#ffffff'):
        cell.setColor('#00ff00')
    else:
        cell.setColor('#ffffff')

stat = 0
tak.GridAgent.step = agentLoop
tak.GridModel.initAgents = agentSetup
tak.Cell.step=cellStep

grassDensity=tak.Var('草密度',0.5,range=(0,0.01,1))# (min,step,max)
sheepGrowth=tak.Var('sheepBirth',0.03,range=(0,0.01,0.05))
grassGrowth=tak.Var('grassGrowth',0.05,range=(0,0.01,0.15))
sheepEnergy=tak.Var('',20)

varList=[grassGrowth,sheepGrowth]
tak.GridModel.varList=varList
tak.GridModel.dataCounters=[tak.DataCounter(name='颜色', propertyName='color', targets=['cell']),
                            tak.DataCounter(propertyName='color', targets=['agent'])]
tak.Cell.onClicked=cellOnClicked
if __name__ == "__main__":
    tak.GridModel.width = 34 # 设置网格宽度
    tak.GridModel.height = 34  # 设置网格高度
    tak.simStart(__file__, tak.GridModel,maxSteps=10000)  # 这个方法同时启动了图形界面。
